import random
import time
import os

import cv2
import logging
import numpy as np

# 导入新的VNC客户端实现
from pytest_vnc import VNCClient
import keyboard


# vnc远程传输虚拟机画面
class VNC:
    # 初始化VNC对象
    def __init__(self, host, port, password=None):
        self.client = None
        self.host = host
        self.port = int(port)
        self.password = password

        self.key_state = {}  # 记录按键状态
        self.current_mouse_pos = (0, 0)  # 初始位置

        # 构建ASCII到VNC按键的映射
        self.ascii_to_vnc = {
            # 数字键
            48: '0', 49: '1', 50: '2', 51: '3', 52: '4',
            53: '5', 54: '6', 55: '7', 56: '8', 57: '9',

            # 字母键(小写)
            65: 'a', 66: 'b', 67: 'c', 68: 'd', 69: 'e',
            70: 'f', 71: 'g', 72: 'h', 73: 'i', 74: 'j',
            75: 'k', 76: 'l', 77: 'm', 78: 'n', 79: 'o',
            80: 'p', 81: 'q', 82: 'r', 83: 's', 84: 't',
            85: 'u', 86: 'v', 87: 'w', 88: 'x', 89: 'y', 90: 'z',

            # 特殊符号
            186: ';', 187: '=', 188: ',', 189: '-', 190: '.',
            191: '/', 192: '`', 219: '[', 220: '\\', 221: ']', 222: "'",

            # 功能键
            8: 'BackSpace', 9: 'Tab', 13: 'Return', 16: 'Shift',
            17: 'Ctrl', 18: 'Alt', 19: 'Pause', 20: 'Caps_Lock',
            27: 'Escape', 32: 'Space', 33: 'Page_Up', 34: 'Page_Down',
            35: 'End', 36: 'Home', 37: 'Left', 38: 'Up', 39: 'Right',
            40: 'Down', 44: 'PrintScreen', 45: 'Insert', 46: 'Delete',

            # 小键盘
            96: 'KP_0', 97: 'KP_1', 98: 'KP_2', 99: 'KP_3', 100: 'KP_4',
            101: 'KP_5', 102: 'KP_6', 103: 'KP_7', 104: 'KP_8', 105: 'KP_9',
            106: 'KP_Multiply', 107: 'KP_Add', 108: 'KP_Enter', 109: 'KP_Subtract',
            110: 'KP_Decimal', 111: 'KP_Divide',

            # F键
            112: 'F1', 113: 'F2', 114: 'F3', 115: 'F4', 116: 'F5',
            117: 'F6', 118: 'F7', 119: 'F8', 120: 'F9', 121: 'F10',
            122: 'F11', 123: 'F12',

            # 锁定键
            144: 'Num_Lock', 145: 'Scroll_Lock',

            # 左右修饰键
            160: 'Shift_L', 161: 'Shift_R', 162: 'Ctrl_L', 163: 'Ctrl_R',
            164: 'Alt_L', 165: 'Alt_R', 91: 'Win_L', 92: 'Win_R'
        }

    # 连接到VNC服务器
    def connect(self, retry=3):
        password_info = "有密码" if self.password else "无密码"
        logging.info(f"VNC正在尝试连接 {self.host}:{self.port} ({password_info})")
        if self.client is not None:
            logging.info(f"VNC已经连接到 {self.host}:{self.port}，无需重复连接")
            return True

        for attempt in range(1, retry + 1):
            try:
                # 使用新的VNC客户端
                logging.info(f"第 {attempt} 次连接尝试...")
                self.client = VNCClient(ip=self.host, port=self.port, timeout=5, password=self.password)
                self.client.connect()

                # 等待客户端准备就绪
                start_time = time.time()
                while not self.client.is_ready() and (time.time() - start_time) < 10:
                    time.sleep(0.1)

                if not self.client.is_ready():
                    raise Exception("VNC客户端在10秒内未准备就绪")

                # 验证vnc是否成功连接（获取一帧数据验证）
                frame = self.capture()
                if frame is None:
                    raise Exception("VNC连接成功但无法获取屏幕数据")

                logging.info(f"VNC已连接到 {self.host}:{self.port}")
                return True
            except Exception as e:
                logging.error(f"VNC第 {attempt} 次连接失败: {e}")
                if self.client:
                    try:
                        self.client.close()
                    except:
                        pass
                    self.client = None
                if attempt < retry:
                    time.sleep(1)

        logging.error(f"VNC经过 {retry} 次尝试，仍无法连接到 {self.host}:{self.port}")
        return False

    # 保存VNC服务器的屏幕截图
    def capture_screen_debug(self, base_path):
        """调试版本的截图方法，保存原始和处理后的图像"""
        if not self.client:
            logging.warning("VNC未连接，无法截图")
            return False

        try:
            # 获取原始帧数据
            raw_frame = self.client.get_frame(copy=True)
            if raw_frame is None:
                logging.warning("获取帧数据失败")
                return False

            # 保存原始数据（未处理）
            if len(raw_frame.shape) == 3 and raw_frame.shape[2] == 4:
                raw_bgr = raw_frame[:, :, :3]  # 去掉alpha通道
            else:
                raw_bgr = raw_frame

            # 保存原始图像
            raw_path = f"{base_path}_raw.png"
            cv2.imwrite(raw_path, raw_bgr, [cv2.IMWRITE_PNG_COMPRESSION, 0])
            logging.info(f"原始图像保存至: {raw_path}")

            # 保存处理后的图像
            processed_frame = self.capture()
            if processed_frame is not None:
                processed_path = f"{base_path}_processed.png"
                cv2.imwrite(processed_path, processed_frame, [cv2.IMWRITE_PNG_COMPRESSION, 0])
                logging.info(f"处理后图像保存至: {processed_path}")
                return True

            return False
        except Exception as e:
            logging.error(f"调试截图失败: {e}")
            return False

    def capture_screen(self, screenshot_path):
        """保存VNC截图，保色清晰化处理"""
        if not self.client:
            logging.warning("VNC未连接，无法截图")
            return False

        try:
            frame = self.capture()
            if frame is not None:
                # 根据文件格式设置保存参数
                if screenshot_path.lower().endswith('.jpg') or screenshot_path.lower().endswith('.jpeg'):
                    # JPEG格式：高质量保存
                    success = cv2.imwrite(screenshot_path, frame, [cv2.IMWRITE_JPEG_QUALITY, 95])
                elif screenshot_path.lower().endswith('.png'):
                    # PNG格式：无损压缩
                    success = cv2.imwrite(screenshot_path, frame, [cv2.IMWRITE_PNG_COMPRESSION, 1])
                else:
                    # 默认保存
                    success = cv2.imwrite(screenshot_path, frame)

                if success:
                    file_size = os.path.getsize(screenshot_path) if os.path.exists(screenshot_path) else 0
                    logging.info(f"保色清晰化截图已保存至 {screenshot_path}, 文件大小: {file_size/1024:.1f}KB")
                    return True
                else:
                    logging.error(f"图像保存失败: {screenshot_path}")
                    return False
            else:
                logging.error("无法获取有效的图像数据")
                return False
        except Exception as e:
            logging.error(f"VNC截图失败: {e}")
            return False

    def capture(self):
        """VNC服务器截图 - 保色清晰化处理（保持原始颜色+提升清晰度）"""
        if not self.client:
            logging.warning("VNC未连接，无法截图")
            return None

        try:
            # 使用新客户端获取帧数据
            frame = self.client.get_frame(copy=True)
            if frame is None:
                logging.warning("获取帧数据失败")
                return None

            logging.info(f"原始帧数据形状: {frame.shape}, 数据类型: {frame.dtype}")

            # 处理颜色通道
            if len(frame.shape) == 3 and frame.shape[2] == 4:
                # 有alpha通道，去掉alpha通道
                frame = frame[:, :, :3]
                logging.info("移除alpha通道")
            elif len(frame.shape) == 3 and frame.shape[2] == 3:
                # 已经是3通道
                logging.info("已经是3通道格式")
            else:
                logging.warning(f"意外的帧格式: {frame.shape}")
                return None

            # 确保数据类型正确
            if frame.dtype != np.uint8:
                frame = frame.astype(np.uint8)

            # 调试：检查颜色值分布
            if frame.size > 0:
                logging.info(f"颜色通道统计 - B通道均值: {np.mean(frame[:,:,0]):.1f}, G通道均值: {np.mean(frame[:,:,1]):.1f}, R通道均值: {np.mean(frame[:,:,2]):.1f}")

            # 保色清晰化处理
            enhanced_frame = self._enhance_image(frame)

            logging.info(f"保色清晰化完成，最终图像形状: {enhanced_frame.shape}")
            return enhanced_frame

        except Exception as e:
            logging.error(f"VNC截图失败: {e}")
            return None

    # 关闭vnc连接
    def close(self):
        if self.client is None:
            logging.warning(f"VNC未建立连接，跳过关闭操作")
            return

        try:
            self.client.close()
            logging.info(f"VNC已安全断开 {self.host}:{self.port}")
        except Exception as e:
            logging.error(f"VNC关闭连接时发生错误: {e}")
        finally:
            self.client = None  # 清理状态

    def _human_delay(self, min_ms=50, max_ms=100):
        """随机延迟模拟人类操作"""
        time.sleep(random.uniform(min_ms / 1000, max_ms / 1000))

    def _enhance_image(self, image):
        """保色清晰化处理 - 在保持原始颜色的基础上提升清晰度"""
        try:
            if image is None or image.size == 0:
                return image

            height, width = image.shape[:2]
            logging.info(f"开始保色清晰化处理，图像尺寸: {width}x{height}")

            # 保存原始图像
            original = image.copy()

            # 1. 智能分辨率提升（如果需要）
            result = original
            if height < 600 or width < 800:
                scale_factor = min(1.5, max(800/width, 600/height))
                if scale_factor > 1.05:
                    new_width = int(width * scale_factor)
                    new_height = int(height * scale_factor)
                    result = cv2.resize(original, (new_width, new_height), interpolation=cv2.INTER_LANCZOS4)
                    logging.info(f"分辨率提升至: {new_width}x{new_height}")

            # 2. 保色降噪 - 使用双边滤波保持边缘和颜色
            # 参数调优：小核，保持颜色和空间相似性
            denoised = cv2.bilateralFilter(result, 5, 25, 25)

            # 3. 保色锐化 - 使用YUV色彩空间只处理亮度
            # 转换到YUV，只对Y通道（亮度）进行处理，保持UV通道（色彩）不变
            yuv = cv2.cvtColor(denoised, cv2.COLOR_BGR2YUV)
            y, u, v = cv2.split(yuv)

            # 对亮度通道进行温和的对比度增强
            clahe = cv2.createCLAHE(clipLimit=1.2, tileGridSize=(8,8))
            y_enhanced = clahe.apply(y)

            # 对亮度通道进行轻微锐化
            y_blurred = cv2.GaussianBlur(y_enhanced, (0, 0), 0.8)
            y_sharpened = cv2.addWeighted(y_enhanced, 1.2, y_blurred, -0.2, 0)
            y_sharpened = np.clip(y_sharpened, 0, 255).astype(np.uint8)

            # 重新组合，U和V通道保持完全不变
            enhanced_yuv = cv2.merge([y_sharpened, u, v])
            result = cv2.cvtColor(enhanced_yuv, cv2.COLOR_YUV2BGR)

            # 4. 最终安全检查 - 确保颜色变化在可接受范围内
            # 计算颜色差异（只关注色彩，不关注亮度）
            original_lab = cv2.cvtColor(original, cv2.COLOR_BGR2LAB)
            result_lab = cv2.cvtColor(result, cv2.COLOR_LAB2BGR)
            result_lab = cv2.cvtColor(result_lab, cv2.COLOR_BGR2LAB)

            # 检查a*b*通道的差异（色彩信息）
            color_diff = np.mean(np.abs(original_lab[:,:,1:] - result_lab[:,:,1:]))

            if color_diff > 5:  # 如果颜色差异太大
                logging.warning(f"颜色变化过大 (差异: {color_diff:.2f})，回退到保守处理")
                # 只进行最基本的锐化，不改变颜色
                kernel = np.array([[0, -0.3, 0],
                                  [-0.3, 2.2, -0.3],
                                  [0, -0.3, 0]])
                result = cv2.filter2D(original, -1, kernel)
                result = np.clip(result, 0, 255).astype(np.uint8)

            # 最终确保数据类型正确
            result = np.clip(result, 0, 255).astype(np.uint8)

            logging.info(f"保色清晰化完成，颜色差异: {color_diff:.2f}")
            return result

        except Exception as e:
            logging.warning(f"保色清晰化处理失败，返回原图: {e}")
            return image

    def _convert_key_to_string(self, key_code):
        """将ASCII键码转换为字符串键名"""
        if key_code in self.ascii_to_vnc:
            return self.ascii_to_vnc[key_code]
        else:
            # 如果没有找到映射，尝试直接使用字符
            try:
                if 32 <= key_code <= 126:  # 可打印ASCII字符
                    return chr(key_code)
            except:
                pass
            raise ValueError(f"Unsupported key code: {key_code}")

    def _bezier_curve(self, start, end, control=None, steps=10):
        """生成贝塞尔曲线路径"""
        if control is None:
            control = (
                (start[0] + end[0]) // 2 + random.randint(-50, 50),  # 减少随机偏移
                (start[1] + end[1]) // 2 + random.randint(-50, 50)
            )

        points = []
        for t in [i / steps for i in range(steps + 1)]:
            x = (1 - t) ** 2 * start[0] + 2 * (1 - t) * t * control[0] + t ** 2 * end[0]
            y = (1 - t) ** 2 * start[1] + 2 * (1 - t) * t * control[1] + t ** 2 * end[1]
            points.append((int(x), int(y)))
        return points

    def key_press(self, key, times=1):
        """按下并弹起按键，支持多次操作"""
        for _ in range(times):
            self.key_down(key)
            self._human_delay(50, 100)  # 按的延迟
            # 10%  概率加多一点延迟，模拟人类思考
            if random.random() < 0.1:
                self._human_delay(20, 40)
            self.key_up(key)
            if times > 1:
                self._human_delay(50, 100)  # 多次按键间额外延迟

    def key_down(self, key):
        """按下按键"""
        if not self.client:
            logging.warning("VNC未连接，无法操作按键")
            return

        try:
            # 先释放相反方向的按键
            if key == keyboard.Right:
                self.key_up(keyboard.Left)
            elif key == keyboard.Left:
                self.key_up(keyboard.Right)
            elif key == keyboard.Up:
                self.key_up(keyboard.Down)
            elif key == keyboard.Down:
                self.key_up(keyboard.Up)

            vnc_key = self._convert_key_to_string(key)
            self.client.key_down(vnc_key)
            logging.info(f"按住：{vnc_key}")
        except Exception as e:
            logging.error(f"按键按下失败: {e}")

    def key_up(self, key):
        """弹起按键"""
        if not self.client:
            logging.warning("VNC未连接，无法操作按键")
            return

        try:
            vnc_key = self._convert_key_to_string(key)
            self.client.key_up(vnc_key)
            logging.info(f"弹起：{vnc_key}")
        except Exception as e:
            logging.error(f"按键弹起失败: {e}")

    def mouse_move(self, x, y):
        """拟人化移动（基于记录的当前位置）"""
        if not self.client:
            logging.warning("VNC未连接，无法操作鼠标")
            return

        try:
            if hasattr(self, 'current_mouse_pos'):
                points = self._bezier_curve(self.current_mouse_pos, (x, y))
                for point in points:
                    self.client.move(*point)
                    time.sleep(random.uniform(0.005, 0.01))  # 加快移动速度
            else:
                self.client.move(x, y)
            self.current_mouse_pos = (x, y)  # 更新位置
        except Exception as e:
            logging.error(f"鼠标移动失败: {e}")

    def mouse_left_click(self, times=1):
        """鼠标左键点击"""
        if not self.client:
            logging.warning("VNC未连接，无法操作鼠标")
            return

        try:
            for _ in range(times):
                x, y = self.current_mouse_pos
                self.client.mouse_down(x, y, 1)
                self._human_delay(150, 400)  # 随机按下延迟
                # 10%  概率加多一点延迟，模拟人类思考
                if random.random() < 0.1:
                    self._human_delay(50, 100)
                self.client.mouse_up(x, y)
                if times > 1:
                    self._human_delay(200, 300)  # 多次点击间延迟
        except Exception as e:
            logging.error(f"鼠标左键点击失败: {e}")

    def mouse_right_click(self, times=1):
        """鼠标右键点击"""
        if not self.client:
            logging.warning("VNC未连接，无法操作鼠标")
            return

        try:
            for _ in range(times):
                x, y = self.current_mouse_pos
                self.client.mouse_down(x, y, 4)  # 右键是按钮掩码4
                self._human_delay(150, 400)  # 随机按下延迟
                # 10%  概率加多一点延迟，模拟人类思考
                if random.random() < 0.1:
                    self._human_delay(50, 100)
                self.client.mouse_up(x, y)
                if times > 1:
                    self._human_delay(200, 300)  # 多次点击间延迟
        except Exception as e:
            logging.error(f"鼠标右键点击失败: {e}")

    def mouse_wheel(self, amount=1):
        """鼠标滚轮滚动"""
        if not self.client:
            logging.warning("VNC未连接，无法操作鼠标")
            return

        try:
            x, y = self.current_mouse_pos
            # 确定滚动方向
            direction = 1 if amount > 0 else -1

            for _ in range(abs(amount)):
                # 使用预先确定的方向
                self.client.scroll_wheel(x, y, direction)
                self._human_delay(30, 60)  # 加快滚轮间隔
                # 5%  概率加多一点延迟，模拟人类思考（降低概率）
                if random.random() < 0.05:
                    self._human_delay(20, 50)
        except Exception as e:
            logging.error(f"鼠标滚轮操作失败: {e}")

    def mouse_left_down(self):
        """鼠标左键按下"""
        if not self.client:
            logging.warning("VNC未连接，无法操作鼠标")
            return

        try:
            x, y = self.current_mouse_pos
            self.client.mouse_down(x, y, 1)
            self._human_delay(50, 100)
            # 10%  概率加多一点延迟，模拟人类思考
            if random.random() < 0.1:
                self._human_delay(50, 100)
        except Exception as e:
            logging.error(f"鼠标左键按下失败: {e}")

    def mouse_left_up(self):
        """鼠标左键弹起"""
        if not self.client:
            logging.warning("VNC未连接，无法操作鼠标")
            return

        try:
            x, y = self.current_mouse_pos
            self.client.mouse_up(x, y)
            self._human_delay(50, 100)
            # 10%  概率加多一点延迟，模拟人类思考
            if random.random() < 0.1:
                self._human_delay(50, 100)
        except Exception as e:
            logging.error(f"鼠标左键弹起失败: {e}")

    def mouse_right_down(self):
        """鼠标右键按下"""
        if not self.client:
            logging.warning("VNC未连接，无法操作鼠标")
            return

        try:
            x, y = self.current_mouse_pos
            self.client.mouse_down(x, y, 4)
            self._human_delay(50, 100)
            # 10%  概率加多一点延迟，模拟人类思考
            if random.random() < 0.1:
                self._human_delay(50, 100)
        except Exception as e:
            logging.error(f"鼠标右键按下失败: {e}")

    def mouse_right_up(self):
        """鼠标右键弹起"""
        if not self.client:
            logging.warning("VNC未连接，无法操作鼠标")
            return

        try:
            x, y = self.current_mouse_pos
            self.client.mouse_up(x, y)
            self._human_delay(50, 100)
            # 10%  概率加多一点延迟，模拟人类思考
            if random.random() < 0.1:
                self._human_delay(50, 100)
        except Exception as e:
            logging.error(f"鼠标右键弹起失败: {e}")

    def role_run(self, key):
        """
        角色奔跑（连续按键操作）
        :param key: 方向键
        """
        if not self.client:
            logging.warning("VNC未连接，无法操作角色移动")
            return

        try:
            # 先释放相反方向的按键
            if key == keyboard.Right:
                self.key_up(keyboard.Left)
            elif key == keyboard.Left:
                self.key_up(keyboard.Right)
            elif key == keyboard.Up:
                self.key_up(keyboard.Down)
            elif key == keyboard.Down:
                self.key_up(keyboard.Up)

            # 按键操作：按下-释放-再按下（保持按压状态）
            self.key_down(key)
            self._human_delay(50, 100)
            self.key_up(key)
            self._human_delay(50, 100)
            self.key_down(key)

        except Exception as e:
            key_name = keyboard.key_name_map.get(key, f"Key_{key}")
            logging.error(f"VNC角色奔跑操作失败 {key_name}: {e}")

    def release_all_key(self):
        """释放所有当前按下的按键"""
        if not self.client:
            logging.warning("VNC未连接，无法释放按键")
            return

        try:
            # 额外确保释放常用的方向键（防止遗漏）
            direction_keys = [keyboard.Q, keyboard.W, keyboard.E, keyboard.R, keyboard.T, keyboard.Y, keyboard.LeftCtrl,
                              keyboard.A, keyboard.S, keyboard.D, keyboard.F, keyboard.G, keyboard.H, keyboard.LeftAlt,
                              keyboard.Up, keyboard.Down, keyboard.Left, keyboard.Right, keyboard.X, keyboard.num_0,
                              keyboard.Esc, keyboard.Space, keyboard.N, keyboard.F1, keyboard.F2, keyboard.F10,
                              keyboard.F12, keyboard.I, keyboard.M, keyboard.J, keyboard.F7, keyboard.P]
            for key in direction_keys:
                self.key_up(key)

        except Exception as e:
            logging.error(f"VNC释放所有按键失败: {e}")

    def simple_key_down(self, key):
        """
        简单按下按键（无延迟检查）
        :param key: 按键码
        """
        if not self.client:
            logging.warning("VNC未连接，无法操作按键")
            return

        try:
            # 先释放相反方向的按键
            if key == keyboard.Right:
                self.key_up(keyboard.Left)
            elif key == keyboard.Left:
                self.key_up(keyboard.Right)
            elif key == keyboard.Up:
                self.key_up(keyboard.Down)
            elif key == keyboard.Down:
                self.key_up(keyboard.Up)

            self.key_down(key)
        except Exception as e:
            key_name = keyboard.key_name_map.get(key, f"Key_{key}")
            logging.error(f"VNC简单按键操作失败 {key_name}: {e}")
